import cv2
import os


def cut_and_save(img_arr, h, w, nh, nw, img_name, save_dir='.'):
    """
    :param img_arr: 待处理的图像数组
    :param nx: 分几行
    :param ny: 分几列
    :param img_name: 原图名字
    :param save_dir: 切图存放目录
    save as png files
    """
    dh, dw = h // nh, w // nw
    for x in range(nh):
        for y in range(nw):
            if x == nh - 1 and y != nw-1:
                small_img = img_arr[x*dh:, y*dw:y*dw+dw, :]
            elif x != nh-1 and y == nw-1:
                small_img = img_arr[x*dh:x*dh+dh, y*dw:, :]
            elif x == nh-1 and y == nw-1:
                small_img = img_arr[x*dh:, y*dw:, :]
            else:
                small_img = img_arr[x*dh:x*dh+dh, y*dw:y*dw+dw, :]
            save_name = f"{img_name}_a_{x}_{y}.png"
            save_path = os.path.join(save_dir, save_name)
            cv2.imwrite(save_path, small_img)


picdir = '.'  # need to be a str
picforms = ['jpg', 'png']  # need to be a list

filenames = os.listdir(picdir)
imglist = [imgname for imgname in filenames if imgname.split('.')[-1] in picforms]
print('len(imglist):', len(imglist))

for img_name in imglist:
    # 读图
    # print(img_name)
    img_arr = cv2.imread(os.path.join(picdir, img_name))
    # print(img_arr.shape)

    # 算尺寸
    h, w = img_arr.shape[0], img_arr.shape[1]
    print(f'{img_name}, size={w}x{h}')

    # 算格子，切几行、切几列
    nh, nw = h // 150 + 1, w // 150 + 1

    # 切图并存图
    save_dir = os.path.join(picdir, img_name.split('.')[-2])
    if not os.path.isdir(save_dir):
        os.mkdir(save_dir)

    cut_and_save(img_arr, h, w, nh, nw, img_name, save_dir)
